import argparse

from vllm_models.showui_model import ShowUIModel


def main(image_path, prompt, draw=False):
    model = ShowUIModel(model_name="AI-ModelScope/ShowUI-2B", base_url="http://10.0.102.69:8000/v1",
                        api_key="123456")
    response = model.process_image(image_path, prompt)
    print(response)
    if draw:
        image = model.draw_point(image_path, response)
        image.save("../examples/labeled_image_0.png")
        print("save image to ../examples/labeled_image_0.png")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Qwen2-VL image description")
    parser.add_argument("--image", required=True, help="Path to the image file")
    parser.add_argument("--prompt", required=True, help="Prompt for the model")
    parser.add_argument("--draw", required=False, action="store_true", help="Draw the image")
    args = parser.parse_args()
    main(args.image, args.prompt, args.draw)
